package com.hgj.common.ws.server;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import jakarta.websocket.OnClose;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
@Component
@EnableScheduling
@ServerEndpoint(value = "/test")
public class HSocket {

    // Redis 连接模板（通过静态方法初始化）
    private static StringRedisTemplate strRedisTemplate;
    private static RedisTemplate<String,Object> redisTemplate;

    // 节点ID（每个实例唯一）
    private static final String NODE_ID = UUID.randomUUID().toString();

    // Redis key前缀
    private static final String SESSION_KEY = "ws:sessions:";
    private static final String SUBSCRIPTION_KEY = "ws:subscriptions:";
    private static final String NODE_SESSIONS_KEY = "ws:node_sessions:" + NODE_ID;
    private static final String DATA_CACHE_KEY = "ws:data_cache:";

    // 本地会话缓存（仅存储当前节点连接）
    private static final Map<String, Session> localSessions = new ConcurrentHashMap<>();

    // 初始化RedisTemplate
    public static void setStrRedisTemplate(StringRedisTemplate stringRedisTemplate) {
        strRedisTemplate= stringRedisTemplate;
    }

    public static void setRedisTemplate(RedisTemplate<String,Object> template) {
        redisTemplate = template;
    }

    public static String getNodeId(){return NODE_ID;}

    public static Map<String, Session> getLocalSessions(){return localSessions;}

    @OnOpen
    public void onOpen(Session session) throws IOException {
        try {
            // 1. 存储会话信息到Redis
            Map<String, String> sessionInfo = new HashMap<>();
            sessionInfo.put("nodeId", NODE_ID);
            sessionInfo.put("createTime", String.valueOf(System.currentTimeMillis()));

            redisTemplate.opsForHash().putAll(SESSION_KEY + session.getId(), sessionInfo);
            redisTemplate.expire(SESSION_KEY + session.getId(), 30, TimeUnit.MINUTES);

            // 2. 添加到节点会话集合
            redisTemplate.opsForSet().add(NODE_SESSIONS_KEY, session.getId());
            redisTemplate.expire(NODE_SESSIONS_KEY, 40, TimeUnit.MINUTES);

            // 3. 存储到本地缓存
            localSessions.put(session.getId(), session);
            log.info("连接成功[{}]: 节点[{}], 当前节点连接数[{}]",
                    session.getId(), NODE_ID, localSessions.size());

        } catch (Exception e) {
            log.error("连接建立失败: {}", e.getMessage());
            session.close();
        }
    }

    @OnClose
    public void onClose(Session session) throws IOException {
        try {
            // 1. 从Redis删除会话信息
            redisTemplate.delete(SESSION_KEY + session.getId());

            // 2. 从节点会话集合移除
            redisTemplate.opsForSet().remove(NODE_SESSIONS_KEY, session.getId());

            // 3. 删除订阅关系
            String subscriptionKey = SUBSCRIPTION_KEY + session.getId();
            Set<Object> types = redisTemplate.opsForSet().members(subscriptionKey);
            if (types != null) {
                types.forEach(type ->
                        redisTemplate.opsForSet().remove(SUBSCRIPTION_KEY + type, session.getId())
                );
            }
            redisTemplate.delete(subscriptionKey);

            // 4. 从本地缓存移除
            localSessions.remove(session.getId());

            log.info("连接关闭[{}]", session.getId());

        } finally {
            session.close();
        }
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        try {
            JSONObject json = JSON.parseObject(message);
            String type = json.getString("type");
            boolean subscribe = json.getBoolean("subscribe");

            // 处理订阅/取消订阅
            handleSubscription(session.getId(), type, subscribe);

            // 立即发送缓存数据
            if (subscribe) {
                sendCachedData(session, type);
            }

        } catch (Exception e) {
            log.error("消息处理失败: {}", e.getMessage());
        }
    }

    private void handleSubscription(String sessionId, String type, boolean subscribe) {
        String sessionSubKey = SUBSCRIPTION_KEY + sessionId;
        String typeSubKey = SUBSCRIPTION_KEY + type;

        if (subscribe) {
            // 添加订阅关系
            redisTemplate.opsForSet().add(sessionSubKey, type);
            redisTemplate.opsForSet().add(typeSubKey, sessionId);
            redisTemplate.expire(sessionSubKey, 40, TimeUnit.MINUTES);
            redisTemplate.expire(typeSubKey, 40, TimeUnit.MINUTES);
        } else {
            // 移除订阅关系
            redisTemplate.opsForSet().remove(sessionSubKey, type);
            redisTemplate.opsForSet().remove(typeSubKey, sessionId);
        }
    }

    private void sendCachedData(Session session, String type) {
        try {
            // 查找类型相关的所有缓存键
            Set<String> keys = redisTemplate.keys(DATA_CACHE_KEY + type + ":*");
            if (keys == null || keys.isEmpty()) return;

            // 构建响应数据
            List<Object> dataList = redisTemplate.opsForValue().multiGet(keys);

            if (dataList == null) return;
            JSONObject response = new JSONObject();
            response.put("type", type);
            response.put("data", JSONArray.parseArray(dataList.toString()));
            session.getBasicRemote().sendText(response.toJSONString());

        } catch (Exception e) {
            log.error("发送缓存数据失败: {}", e.getMessage());
        }
    }

    // 发布消息到所有订阅者（分布式广播）
    public static void publishMessage(String type, JSONObject data, String keyField) {
        try {
            String keyValue = data.getString(keyField);
            String cacheKey = DATA_CACHE_KEY + type;
            if(keyValue!=null) cacheKey += ":" + keyValue;
            //判断数据是否重复
            String dataStr = data.toJSONString();
            Object cacheData = redisTemplate.opsForValue().get(cacheKey);
            if(cacheData==null || !cacheData.equals(dataStr)){
                // 1. 存入/更新缓存
                redisTemplate.opsForValue().set(cacheKey, dataStr, 24, TimeUnit.HOURS);
                // 2. 构建广播消息
                JSONObject message = new JSONObject();
                message.put("type", type);
                message.put("data", keyValue!=null?Collections.singletonList(data):data);
                String messageStr = message.toJSONString();
                // 3. 获取所有订阅该类型的会话
                Set<Object> sessionIds = redisTemplate.opsForSet().members(SUBSCRIPTION_KEY + type);
                if (sessionIds == null) return;
                // 4. 按节点分组发送
                Map<String, Set<String>> nodeSessions = new HashMap<>();
                for (Object sessionId : sessionIds) {
                    // 获取会话所在节点
                    String nodeId = (String) redisTemplate.opsForHash()
                            .get(SESSION_KEY + sessionId, "nodeId");

                    if (nodeId != null) {
                        nodeSessions.computeIfAbsent(nodeId, k -> new HashSet<>())
                                .add(sessionId.toString());
                    }
                }
                // 5. 向各节点发送广播任务
                nodeSessions.forEach((nodeId, ids) -> {
                    if (NODE_ID.equals(nodeId)) {
                        // 当前节点直接处理
                        sendToLocalSessions(ids, messageStr);
                    } else {
                        // 其他节点通过Redis队列通知
                        redisTemplate.convertAndSend("ws:broadcast:" + nodeId,
                                JSON.toJSONString(new BroadcastTask(ids, messageStr))
                        );
                    }
                });
            }
        } catch (Exception e) {
            log.error("消息发布失败: {}", e.getMessage());
        }
    }


    // 发送消息到本地会话
    public static void sendToLocalSessions(Set<String> sessionIds, String message) {
        sessionIds.forEach(sessionId -> {
            Session session = localSessions.get(sessionId);
            if (session != null && session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (IOException e) {
                    log.error("本地消息发送失败: {}", e.getMessage());
                }
            }
        });
    }

    // 获取全局连接数
    public int getGlobalConnectionCount() {
        Set<String> nodeKeys = redisTemplate.keys("ws:node_sessions:*");
        if (nodeKeys != null) {
            return nodeKeys.stream()
                    .mapToInt(key -> redisTemplate.opsForSet().size(key).intValue())
                    .sum();
        }
        return 0;
    }

    //// 定时延长TTL
    @Scheduled(fixedRate = 15 * 60 * 1000)
    public void renewSessionTTL() {
        HSocket.getLocalSessions().keySet().forEach(sessionId -> {
            redisTemplate.expire(SESSION_KEY + sessionId, 30, TimeUnit.MINUTES);
            redisTemplate.expire(SUBSCRIPTION_KEY + sessionId, 40, TimeUnit.MINUTES);
        });
        redisTemplate.expire(NODE_SESSIONS_KEY, 40, TimeUnit.MINUTES);
    }

    // 广播任务数据结构
    public static class BroadcastTask {
        public BroadcastTask(){};
        public Set<String> sessionIds;
        public String message;

        BroadcastTask(Set<String> sessionIds, String message) {
            this.sessionIds = sessionIds;
            this.message = message;
        }
    }
}